(defpackage :tree-to-equation
  (:use :cl)
  (:nicknames :t2e)
  (:export :tree-to-equation))

(in-package :tree-to-equation)

(defgeneric tree-to-equation-method (root variables))

(defun var-structure ()
  (list :vars))

(defstruct (equation-variable (:conc-name var-))
  symbol node value)

(defun new-var (vars node value)
  "Use side effects to add the variable to the list."
  (let ((v (or (loop for i in '(x y z a b c d e f g h i j k l m n o p q r s t u v w)
                  unless (find i (cdr vars) :key 'var-symbol :test 'string=)
                  return i)
               (error "there are too many X boxes!"))))
    ;; red unknowns get a special variable name
    (when (and (typep node 'd:box)
               (eql (length (tree:children node)) 1)
               (sdl:color= (d:background (tree:child node)) (gwee:color 255 50 50)))
      (setf v (intern (format nil "~ar" v) (symbol-package v))))
    (nconc vars (list (make-equation-variable :symbol v :node node :value value)))
    v))

(defmethod tree-to-equation-method ((root d:multi-argument-box) vars)
  (ecase (length (tree:children root))
    (1 (if (tree:children (tree:child root))
           (list '= (tree-to-equation-method (tree:child root) vars) "?")
           (new-var vars root (tree-to-equation-method (tree:child root) vars))))
    (3
     (list (d:quantity (second (tree:children root)))
           (tree-to-equation-method (first (tree:children root)) vars)
           (tree-to-equation-method (third (tree:children root)) vars)))))

(defmethod tree-to-equation-method :around ((root d:multi-argument-box) vars)
  (let ((var (find root (cdr vars) :key 'var-node)))
    (cond
      (var
       (var-symbol var))
      ((and (> (length (tree:parents root)) 1)
            (not (d:xbox-p root)))
         (new-var vars root (call-next-method)))
      (t
       (call-next-method)))))

(defun ditch-extraneous-ops (tree)
  (if (and (listp tree)
           (or (and
                (eq (first tree) '=)
                (equal (third tree) "?"))
               (and
                (eq (first tree) '\|)
                (eql (length tree) 2))))
      (second tree)
      tree))

(defmethod tree-to-equation-method ((root d::invisible-box) vars)
  (labels ((do-child (n)
             (ditch-extraneous-ops (tree-to-equation-method n vars))))
    (cons '\| (mapcar #'do-child (tree:children root)))))

(defmethod tree-to-equation-method ((root d:argument) vars)
  (cond
    ((and (tree:children root)
          (d:quantity root))
     `(= ,@(mapcar (lambda (x) (tree-to-equation-method x vars)) (tree:children root))
         ,(d:quantity root)))
    ((> (length (tree:children root)) 1)
     `(= ,@(mapcar (lambda (x) (tree-to-equation-method x vars)) (tree:children root))))
    ((tree:children root)
     (tree-to-equation-method (tree:child root) vars))
    (t (or (d:quantity root)
           "?"))))

(defun simplify-tree (equation)
  (labels ((helper (e)
             (if (listp e)
                 (let ((elts (mapcar #'helper (cdr e))))
                   (if (and (listp (car elts))
                            (eq (caar elts) (car e)))
                       `(,(car e) ,@(cdar elts) ,@(cdr elts))
                       `(,(car e) ,@elts)))
                 e)))
    (helper equation)))

(defun infix-tree (equation omit-*)
  (labels ((higher-precedence-p (op1 op2)
             (cond
               ((null op1)
                nil)
               ((eq op1 '\|)
                (eq op2 '\|))
               ((eq op1 '=)
                (find op2 '(= \|)))
               ((find op1 '(+ -))
                (find op2 '(+ - = \|)))
               ((find op1 '(* /))
                (find op2 '(\| = + - * /)))
               (t (error "Dunno how to handle op ~s" op1))))
           (helper (e parent-op)
             (if (listp e)
                 (let* ((op (first e))
                        (parens (higher-precedence-p parent-op op)))
                   (format nil (format nil "~a~~{~~a~~^~a~~}~a"
                                       (if parens "(" "")
                                       (if (and omit-*
                                                (eq op '*)
                                                (numberp (second e))
                                                (symbolp (third e)))
                                           "" op)
                                       (if parens ")" ""))
                           (mapcar (u:curry helper (x) op) (cdr e))))
                 (format nil "~a" e))))
    (helper (simplify-tree equation) nil)))

(defun infix-equation (tree vars omit-*)
  (format nil "~a~{,~a~}"
          (infix-tree tree omit-*)
          (mapcar (lambda (v) (format nil "~a=~a" (second v) (infix-tree (third v) omit-*)))
                  vars)))

(defun tree-to-equation (root &optional omit-*)
  "OMIT-* will write N*X as NX, eg 17*Y -> 17Y. It's for readibility."
  (let ((vars (var-structure)))
    (infix-equation
     (ditch-extraneous-ops (tree-to-equation-method root vars))
     (loop for v in (cdr vars)
        when (and (var-value v)
                  (u:nequal (var-value v) "?"))
        collect `(= ,(var-symbol v) ,(var-value v)))
     omit-*)))
